"""
Neo4j图谱存储实现
支持复杂的关系查询和图谱分析
"""

import asyncio
from typing import List, Dict, Any, Optional
from neo4j import AsyncGraphDatabase, AsyncDriver, AsyncSession
import json
from datetime import datetime

from ...core.interfaces import GraphStoreInterface, Entity, Relationship
from ...core.events import EventEmitter
from ...core.config import config


class Neo4jGraphStore(GraphStoreInterface, EventEmitter):
    """Neo4j图谱存储实现"""
    
    def __init__(
        self, 
        uri: str = "bolt://localhost:7687",
        username: str = "neo4j", 
        password: str = "password",
        database: str = "neo4j"
    ):
        super().__init__()
        self.uri = uri
        self.username = username
        self.password = password
        self.database = database
        self.driver: Optional[AsyncDriver] = None
        
    async def initialize(self):
        """初始化Neo4j连接"""
        try:
            self.driver = AsyncGraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password)
            )
            
            # 测试连接
            await self.driver.verify_connectivity()
            
            # 设置数据库模式
            await self.setup_schema()
            
            await self.emit_success("graph_store_initialized", {
                "uri": self.uri,
                "database": self.database
            })
            
        except Exception as e:
            await self.emit_error("graph_store_initialization", e)
            raise
    
    async def setup_schema(self):
        """设置图谱模式（索引和约束）"""
        schema_queries = [
            # 创建唯一约束
            "CREATE CONSTRAINT person_id IF NOT EXISTS FOR (p:Person) REQUIRE p.id IS UNIQUE",
            "CREATE CONSTRAINT company_id IF NOT EXISTS FOR (c:Company) REQUIRE c.id IS UNIQUE", 
            "CREATE CONSTRAINT skill_id IF NOT EXISTS FOR (s:Skill) REQUIRE s.id IS UNIQUE",
            "CREATE CONSTRAINT position_id IF NOT EXISTS FOR (pos:Position) REQUIRE pos.id IS UNIQUE",
            "CREATE CONSTRAINT education_id IF NOT EXISTS FOR (e:Education) REQUIRE e.id IS UNIQUE",
            
            # 创建索引
            "CREATE INDEX person_name IF NOT EXISTS FOR (p:Person) ON (p.name)",
            "CREATE INDEX company_name IF NOT EXISTS FOR (c:Company) ON (c.name)",
            "CREATE INDEX skill_name IF NOT EXISTS FOR (s:Skill) ON (s.name)",
            "CREATE INDEX position_title IF NOT EXISTS FOR (pos:Position) ON (pos.title)",
        ]
        
        async with self.driver.session(database=self.database) as session:
            for query in schema_queries:
                try:
                    await session.run(query)
                except Exception as e:
                    # 约束已存在时会抛出异常，忽略
                    pass
    
    async def add_entity(self, entity: Entity) -> bool:
        """添加实体"""
        if not self.driver:
            await self.initialize()
        
        try:
            # 动态构建Cypher查询
            properties_str = ", ".join([f"{k}: ${k}" for k in entity.properties.keys()])
            query = f"""
            MERGE (e:{entity.type} {{id: $id}})
            SET e += {{{properties_str}}}
            SET e.updated_at = datetime()
            RETURN e
            """
            
            params = {
                "id": entity.id,
                **entity.properties
            }
            
            async with self.driver.session(database=self.database) as session:
                result = await session.run(query, params)
                record = await result.single()
                
                if record:
                    await self.emit("entity_added", {
                        "entity_id": entity.id,
                        "entity_type": entity.type
                    })
                    return True
                    
            return False
            
        except Exception as e:
            await self.emit_error("add_entity", e)
            return False
    
    async def add_relationship(self, relationship: Relationship) -> bool:
        """添加关系"""
        if not self.driver:
            await self.initialize()
        
        try:
            # 构建关系属性
            props_str = ""
            if relationship.properties:
                prop_pairs = [f"{k}: ${k}" for k in relationship.properties.keys()]
                props_str = f" {{{', '.join(prop_pairs)}}}"
            
            query = f"""
            MATCH (source {{id: $source_id}})
            MATCH (target {{id: $target_id}})
            MERGE (source)-[r:{relationship.type}{props_str}]->(target)
            SET r.updated_at = datetime()
            RETURN r
            """
            
            params = {
                "source_id": relationship.source,
                "target_id": relationship.target,
                **relationship.properties
            }
            
            async with self.driver.session(database=self.database) as session:
                result = await session.run(query, params)
                record = await result.single()
                
                if record:
                    await self.emit("relationship_added", {
                        "source": relationship.source,
                        "target": relationship.target,
                        "type": relationship.type
                    })
                    return True
                    
            return False
            
        except Exception as e:
            await self.emit_error("add_relationship", e)
            return False
    
    async def find_related_entities(
        self, 
        entity_id: str, 
        relation_types: Optional[List[str]] = None,
        max_depth: int = 2
    ) -> List[Entity]:
        """查找相关实体"""
        if not self.driver:
            await self.initialize()
        
        try:
            # 构建关系类型过滤
            rel_filter = ""
            if relation_types:
                rel_types_str = "|".join(relation_types)
                rel_filter = f":{rel_types_str}"
            
            query = f"""
            MATCH (start {{id: $entity_id}})
            MATCH (start)-[r{rel_filter}*1..{max_depth}]-(related)
            WHERE related.id <> $entity_id
            RETURN DISTINCT related, labels(related) as types
            LIMIT 100
            """
            
            async with self.driver.session(database=self.database) as session:
                result = await session.run(query, {"entity_id": entity_id})
                
                entities = []
                async for record in result:
                    node = record["related"]
                    node_types = record["types"]
                    
                    # 取第一个标签作为主要类型
                    entity_type = node_types[0] if node_types else "Unknown"
                    
                    # 转换节点属性
                    properties = dict(node)
                    
                    entity = Entity(
                        id=properties.pop("id"),
                        type=entity_type,
                        properties=properties
                    )
                    entities.append(entity)
                
                await self.emit("related_entities_found", {
                    "entity_id": entity_id,
                    "count": len(entities),
                    "max_depth": max_depth
                })
                
                return entities
                
        except Exception as e:
            await self.emit_error("find_related_entities", e)
            return []
    
    async def query_graph(self, cypher_query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
        """执行自定义Cypher查询"""
        if not self.driver:
            await self.initialize()
        
        try:
            async with self.driver.session(database=self.database) as session:
                result = await session.run(cypher_query, parameters or {})
                
                records = []
                async for record in result:
                    # 转换记录为字典
                    record_dict = {}
                    for key in record.keys():
                        value = record[key]
                        # 处理Neo4j节点和关系对象
                        if hasattr(value, '__dict__'):
                            record_dict[key] = dict(value)
                        else:
                            record_dict[key] = value
                    records.append(record_dict)
                
                await self.emit("graph_query_executed", {
                    "query": cypher_query[:100] + "..." if len(cypher_query) > 100 else cypher_query,
                    "results_count": len(records)
                })
                
                return records
                
        except Exception as e:
            await self.emit_error("query_graph", e)
            return []
    
    async def find_shortest_path(self, source_id: str, target_id: str, max_length: int = 6) -> Optional[List[Dict[str, Any]]]:
        """查找两个实体间的最短路径"""
        query = """
        MATCH (source {id: $source_id}), (target {id: $target_id})
        MATCH path = shortestPath((source)-[*1..{max_length}]-(target))
        RETURN nodes(path) as nodes, relationships(path) as relationships
        """.format(max_length=max_length)
        
        result = await self.query_graph(query, {
            "source_id": source_id,
            "target_id": target_id
        })
        
        return result[0] if result else None
    
    async def get_entity_neighbors(self, entity_id: str, limit: int = 50) -> Dict[str, List[Entity]]:
        """获取实体的直接邻居，按关系类型分组"""
        query = """
        MATCH (center {id: $entity_id})-[r]-(neighbor)
        RETURN type(r) as relation_type, collect(DISTINCT neighbor) as neighbors
        LIMIT $limit
        """
        
        result = await self.query_graph(query, {
            "entity_id": entity_id,
            "limit": limit
        })
        
        neighbors_by_type = {}
        for record in result:
            relation_type = record["relation_type"]
            neighbor_nodes = record["neighbors"]
            
            entities = []
            for node in neighbor_nodes:
                entity = Entity(
                    id=node.get("id"),
                    type=list(node.labels)[0] if node.labels else "Unknown",
                    properties=dict(node)
                )
                entities.append(entity)
            
            neighbors_by_type[relation_type] = entities
        
        return neighbors_by_type
    
    async def get_graph_stats(self) -> Dict[str, Any]:
        """获取图谱统计信息"""
        stats_queries = {
            "total_nodes": "MATCH (n) RETURN count(n) as count",
            "total_relationships": "MATCH ()-[r]->() RETURN count(r) as count",
            "node_types": "MATCH (n) RETURN labels(n) as labels, count(n) as count",
            "relationship_types": "MATCH ()-[r]->() RETURN type(r) as type, count(r) as count"
        }
        
        stats = {}
        for stat_name, query in stats_queries.items():
            try:
                result = await self.query_graph(query)
                if stat_name in ["total_nodes", "total_relationships"]:
                    stats[stat_name] = result[0]["count"] if result else 0
                else:
                    stats[stat_name] = result
            except Exception:
                stats[stat_name] = 0 if "total" in stat_name else []
        
        return stats
    
    async def close(self):
        """关闭数据库连接"""
        if self.driver:
            await self.driver.close()
            await self.emit("graph_store_closed", {})